import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

#欧氏距离函数
def distance(a, b):
    """计算两个样本点之间的欧氏距离"""
    return np.sqrt(np.sum((a - b) ** 2))

#KNN算法类
class KNN:
    def __init__(self, k, label_num):
        self.k = k
        self.label_num = label_num  # 类别数量

    def fit(self, x_train, y_train):
        """保存训练数据"""
        self.x_train = np.array(x_train)
        self.y_train = np.array(y_train)

    def get_knn_indices(self, x):
        """计算样本 x 到所有训练样本的距离，并返回最近 k 个索引"""
        dis = list(map(lambda a: distance(a, x), self.x_train))
        knn_indices = np.argsort(dis)[:self.k]
        return knn_indices

    def get_label(self, x):
        """根据K个近邻的标签进行投票"""
        knn_indices = self.get_knn_indices(x)
        label_statistic = np.zeros(shape=[self.label_num])
        for index in knn_indices:
            label = int(self.y_train[index])
            label_statistic[label] += 1
        return np.argmax(label_statistic)

    def predict(self, x_test):
        """预测测试集样本的类别"""
        predicted_labels = np.zeros(shape=[len(x_test)], dtype=int)
        for i, x in enumerate(x_test):
            predicted_labels[i] = self.get_label(x)
        return predicted_labels

if __name__ == "__main__":
    # 模拟训练集（二维特征，两类）
    x_train = np.array([
        [1, 2], [2, 3], [3, 3],
        [6, 5], [7, 7], [8, 6]
    ])
    y_train = np.array([0, 0, 0, 1, 1, 1])

    # 模拟测试集
    x_test = np.array([
        [2, 2], [7, 5], [5, 4]
    ])
    y_test = np.array([0, 1, 1])  # 真实标签

    # 选择K并预测
    k = 3
    knn = KNN(k, label_num=2)
    knn.fit(x_train, y_train)
    y_pred = knn.predict(x_test)
    acc = np.mean(y_pred == y_test)
    print(f"K = {k}, 预测结果 = {y_pred}, 准确率 = {acc * 100:.1f}%")

    #绘制可视化
    plt.figure(figsize=(6, 5))
    plt.title(f"KNN分类结果 (k={k}, 准确率={acc*100:.1f}%)")

    # 绘制训练数据点
    plt.scatter(x_train[y_train == 0][:, 0], x_train[y_train == 0][:, 1],
                color='blue', label='类别0 (训练)')
    plt.scatter(x_train[y_train == 1][:, 0], x_train[y_train == 1][:, 1],
                color='red', label='类别1 (训练)')

    # 绘制测试数据点（星形标记）
    for i in range(len(x_test)):
        plt.scatter(x_test[i, 0], x_test[i, 1],
                    color='green' if y_pred[i] == 0 else 'orange',
                    marker='*', s=180,
                    edgecolor='black',
                    label=f"测试点{i+1} (预测{y_pred[i]})" if i < 2 else None)

    plt.legend()
    plt.xlabel('X1')
    plt.ylabel('X2')
    plt.grid(True)
    plt.show()
